package com.xw.echo.core.trans;

import com.xw.echo.core.ContainerMap;
import com.xw.echo.core.pool.ConnectPool;
import com.xw.echo.core.protocol.MethodInvokData;
import io.netty.channel.Channel;
import io.netty.channel.EventLoop;
import io.netty.util.concurrent.DefaultPromise;
import io.netty.util.concurrent.Promise;
import io.netty.util.concurrent.ScheduledFuture;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
public class EchoTransPort implements TransPort {
    private ConnectPool connectPool;
    private Map<String, Promise<Object>> responsePromises = ContainerMap.PromiseMap;


    public EchoTransPort(ConnectPool connectPool) {
        this.connectPool = connectPool;
    }

    @Override
    public Object invoker(MethodInvokData methodInvokData) throws InterruptedException, ExecutionException {
        methodInvokData.setRequestId(generateId());
        Channel channel = getChannel();
        Promise<Object> resultPromise = getPromise(channel, methodInvokData);
        Object result=null;
        channel.writeAndFlush(methodInvokData).addListener(future -> {
            if (!future.isSuccess()) {
                resultPromise.setFailure(future.cause());
            }
        });
        if (resultPromise.await(8, TimeUnit.SECONDS)) {
            if (resultPromise.isSuccess()) {
                result= resultPromise.get();
            } else {
                // 处理请求失败的情况
                /*  throw new RuntimeException("Request failed: " + resultPromise.cause());*/
                responsePromises.remove(methodInvokData.getRequestId());
                log.error("Request failed: {}", resultPromise.cause().getMessage());
                throw new RuntimeException("Request failed: " + resultPromise.cause().getMessage());
            }
        } else {
            // 超时，移除 Promise 并抛出异常
            responsePromises.remove(methodInvokData.getRequestId());
            throw new RuntimeException("RPC Call time out");
        }
        if(channel.isActive())
        {
            connectPool.returnChannel(channel);
        }
        return result;
    }


    @Override
    public CompletableFuture<Object> asyncInvoker(MethodInvokData methodInvokData) {
        // 在这里设置 Request ID!
        methodInvokData.setRequestId(generateId());

        CompletableFuture<Channel> channelFuture = CompletableFuture.supplyAsync(this::getChannel);
        return channelFuture.thenCompose(channel -> {
            // 确保在继续之前 channel 不为 null (之前的修复)
            if (channel == null) {
                CompletableFuture<Object> failedFuture = new CompletableFuture<>();
                failedFuture.completeExceptionally(new RuntimeException("未能获取用于异步调用的 Channel"));
                return failedFuture;
            }

            // 现在调用 getPromise 时，methodInvokData 已经有 requestId 了
            Promise<Object> resultPromise = getPromise(channel, methodInvokData);
            CompletableFuture<Object> resultFuture = new CompletableFuture<>();
            EventLoop eventLoop = channel.eventLoop();
            AtomicBoolean released = new AtomicBoolean(false); // 标记只释放一次 (之前的修复)
            Runnable releaseAction = () -> {
                if (released.compareAndSet(false, true)) {
                    connectPool.releaseChannel(channel);
                }
            };

            // --- 后续代码保持不变 ---

            // 异步写入数据
            channel.writeAndFlush(methodInvokData).addListener(future -> {
                if (!future.isSuccess()) {
                    resultPromise.tryFailure(future.cause());
                    resultFuture.completeExceptionally(future.cause());
                    releaseAction.run();
                }
            });

            // 处理 Promise 完成
            resultPromise.addListener(promise -> {
                if (promise.isSuccess()) {
                    resultFuture.complete(promise.get());
                } else {
                    resultFuture.completeExceptionally(promise.cause());
                }
                releaseAction.run();
            });

            // 安排超时任务
            ScheduledFuture<?> timeoutTask = eventLoop.schedule(() -> {
                if (!resultPromise.isDone()) {
                    TimeoutException timeoutException = new TimeoutException("Operation timed out after 8 seconds");
                    if (resultPromise.tryFailure(timeoutException)) {
                        resultFuture.completeExceptionally(timeoutException);
                        releaseAction.run();
                    }
                }
            }, 8, TimeUnit.SECONDS);

            // 如果 Future 提前完成取消超时任务
            resultFuture.whenComplete((res, ex) -> {
                if (!timeoutTask.isDone()) {
                    timeoutTask.cancel(false);
                }
            });

            return resultFuture;
        });
    }

private Channel getChannel() {
    Channel channel = connectPool.getChannel();
    if (channel == null || !channel.isActive()) {
        // 处理获取连接失败的情况
        throw new RuntimeException("Failed to get a channel");
    } else {
        return channel;
    }

}
    private String generateId() {
        return UUID.randomUUID().toString();
    }
private Promise<Object> getPromise(Channel channel, MethodInvokData methodInvokData) {
    String requestId = methodInvokData.getRequestId();

    Promise<Object> resultPromise = new DefaultPromise<>(channel.eventLoop());
    responsePromises.put(requestId, resultPromise);
    return resultPromise;
}
}
